package P4树形结构;

public class Trie {
    private class Node {
        public boolean isWord;  //到此为止 是否成立一个单词
        public Map<Character, Node> nexts;   //该字母后续下一层所有的字母
        public int count;   //到此为止 成立单词的个数 记录重复

        public Node() {
            isWord = false;
            count = 0;
            nexts = new AVLTreeMap<>();
        }
    }

    private Node root;
    private int size;   //不重复元素的个数
    private int count;  //所有元素的个数(包含重复)

    public Trie() {
        root = new Node();
        size = 0;
        count = 0;
    }

    public void add(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.nexts.contains(c)) {
                cur.nexts.put(c, new Node());
            }
            cur = cur.nexts.get(c);
        }
        if (!cur.isWord) {
            cur.isWord = true;
            size++;
        }
        cur.count++;
        count++;
    }

    public boolean contains(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.nexts.contains(c)) {
                return false;
            }
            cur = cur.nexts.get(c);
        }
        return cur.isWord;
    }

    public int wordCount(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.nexts.contains(c)) {
                return 0;
            }
            cur = cur.nexts.get(c);
        }
        if (cur.isWord) {
            return cur.count;
        }
        return 0;
    }

    public boolean isPrefix(String prefix) {
        Node cur = root;
        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (!cur.nexts.contains(c)) {
                return false;
            }
            cur = cur.nexts.get(c);
        }
        return true;
    }

    public int size() {
        return size;
    }
    public int count() {
        return count;
    }
}